import ast
import sys
import os
from jinja2 import Template
import json
import neptune
import re
from collections import defaultdict

name_filter = sys.argv[1]

def string_after(hay, needle):
  if needle in hay:
    return hay[hay.find(needle) + len(needle):]
  return None

def raw_config_to_modern(params, config, name='original'):
  params['lotsize'][name] = config['lot_size']
  params['horizon'][name] = re.split(r'_|\.', config['linear_alpha_file'])[-2][:-1]
  params['threshold'][name] = config['threshold_bps']
  params['stack'][name] = (config['max_pos'] - config['min_pos']) / config['lot_size']
  params['lean_per_stack_bps'][name] = config.get('lean_per_stack_bps', '')
  params['pie_in_the_sky_cap_bps'][name] = config.get('pie_in_the_sky_cap_bps', '')
  params['cancel_edge_bps'][name] = config.get('cancel_edge_bps', '')
  params['cancel_outside_bps'][name] = config.get('cancel_outside_bps', '')

  for key in ['use_passive_post', 'pie_in_the_sky_cap_bps']:
    params[key][name] = config.get(key, '?')
  params['stack'][name] = config['max_pos'] / config['lot_size'] * 2


session = neptune.sessions.Session.with_default_backend().get_project('donggu-prestolabs/lm-agg2')
# exps = session.get_experiments(state=['succeeded', 'running', 'failed'])
exps = session.get_experiments(tag=[name_filter])

reports = []
os.makedirs('plots', exist_ok=True)
for exp in exps:
  print('processing ', exp.id, exp.name)
  if exp.name != name_filter:
    continue

  logs = exp.get_logs()
  prop = exp.get_properties()

  # until we impl
  # exp.download_artifact(f'Trial_{best_trial}.png', 'plots')
  # n_trials = prop['n_trials']

  exchange = prop['exchange']
  symbol = prop['target']
  exp_dir = f'out/{exp.name}-{exchange}-{symbol}'
  stripped_exp_dir = f'{exp.name}-{exchange}-{symbol}'

  # until we impl
  # best_trial = prop['best_trial']
  try:
    df = exp.get_numeric_channels_values('run_score')
  except Exception:
    print('no run_score field. skip.')
    continue
  best_trial = df.run_score.idxmax()
  best_score = df.run_score[best_trial]
  original_score = df.run_score[0]
  n_trials = len(df)

  # download config and read original / best params
  params = defaultdict(dict)
  exp.download_artifact('configs/Trial_0.json', exp_dir)
  exp.download_artifact(f'configs/Trial_{best_trial}.json', exp_dir)
  config = json.loads(open(f'{exp_dir}/Trial_0.json').read())['products'][symbol]
  raw_config_to_modern(params, config, 'original')
  config = json.loads(open(f'{exp_dir}/Trial_{best_trial}.json').read())['products'][symbol]
  raw_config_to_modern(params, config, 'proposed')

  if exp.state == 'succeeded' or exp.state == 'failed':
    for t in [0, best_trial]:
      if not os.path.exists(f'{exp_dir}/sim_trial{t}_total.png'):
        try:
          exp.download_artifact(f'pnls/Trial_{t}.png', exp_dir)
          os.rename(
              f'{exp_dir}/Trial_{t}.png',
              f'{exp_dir}/sim_trial{t}_total.png')
        except Exception as e:
          print('pnl plot not found', t, e)
          pass
    if not os.path.exists(f'{exp_dir}/slice.html'):
      try:
        exp.download_artifact('charts/slice.html', exp_dir)
      except Exception:
        pass
    if not os.path.exists(f'{exp_dir}/parallel_coordinate.html'):
      try:
        exp.download_artifact('charts/parallel_coordinate.html', exp_dir)
      except Exception:
        pass
    if not os.path.exists(f'{exp_dir}/param_importances.html'):
      try:
        exp.download_artifact('charts/param_importances.html', exp_dir)
      except:
        pass

  r = dict(
      n_trials=n_trials,
      best_trial=best_trial,
      original_score=original_score,
      best_score=best_score,
      original_config_link=f'{stripped_exp_dir}/Trial_0.json',
      proposal_config_link=f'{stripped_exp_dir}/Trial_{best_trial}.json',
      exchange=prop['exchange'],
      symbol=prop['target'],
      params=params,
      dates=prop['dates'],
      original_plot=f'{stripped_exp_dir}/sim_trial0_total.png',
      proposed_plot=f'{stripped_exp_dir}/sim_trial{best_trial}_total.png',
      slice_plot=f'{stripped_exp_dir}/slice.html',
      parallel_plot=f'{stripped_exp_dir}/parallel_coordinate.html',
      param_importance_plot=f'{stripped_exp_dir}/param_importances.html'
  )
  reports.append(r)

f = open('cc/appcoin2/experimental/donggu/research/report/index.jinja2')
template = Template(f.read())

f2 = open(f'out/report_{name_filter}.html', 'w')
f2.write(template.render(reports=reports))

print(f'out/report_{name_filter}.html')
